# words/management/commands/import_words.py

import os
import pandas as pd
from django.core.management.base import BaseCommand
from django.db import transaction
from words.models import Noun, Verb, Adjective, Other, VerbConjugation

class Command(BaseCommand):
    help = 'Import words from Excel files into the database'

    def add_arguments(self, parser):
        parser.add_argument(
            '--excel_dir',
            type=str,
            help='Path to the directory containing Excel files',
            default='excel_files'  # 默认目录
        )

    def handle(self, *args, **kwargs):
        excel_dir = kwargs['excel_dir']
        base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        excel_path = os.path.join(base_dir, excel_dir)

        if not os.path.exists(excel_path):
            self.stdout.write(self.style.ERROR(f'Excel文件目录不存在: {excel_path}'))
            return

        # 文件名与模型的对应关系
        files = {
            '名词.xlsx': Noun,
            '动词.xlsx': Verb,
            '形容词.xlsx': Adjective,
            '其他.xlsx': Other,
        }

        for file_name, model in files.items():
            file_path = os.path.join(excel_path, file_name)
            if not os.path.exists(file_path):
                self.stdout.write(self.style.ERROR(f'文件不存在: {file_path}'))
                continue

            self.stdout.write(self.style.NOTICE(f'正在导入文件: {file_path}'))

            try:
                # 读取Excel文件
                df = pd.read_excel(file_path)

                with transaction.atomic():
                    if model == Verb:
                        # 清空现有数据
                        model.objects.all().delete()
                        VerbConjugation.objects.all().delete()
                        self.stdout.write(f'已清空 {model.__name__} 表和 VerbConjugation 表中的数据')

                        # 将DataFrame转换为字典列表
                        records = df.to_dict(orient='records')

                        # 处理布尔字段，以及动词类别verb_class
                        for record in records:
                            record['stark'] = bool(record['stark'])
                            record['divisible'] = bool(record['divisible'])
                            record['reflexive'] = bool(record['reflexive'])
                            record['auxiliary'] = bool(record['auxiliary'])

                            # 将 Excel 中的 "class" 字段转换为模型中的 "verb_class" 字段
                            # 假设 Excel 中的列名就是 "class"
                            # 如果 Excel 列名与此不一致，请做相应修改
                            excel_class_value = record.pop('class', 0)  # 若没有该列，默认为0
                            record['verb_class'] = int(excel_class_value)

                        # 批量创建 Verb 对象
                        verb_objs = [model(**record) for record in records]
                        model.objects.bulk_create(verb_objs)
                        self.stdout.write(self.style.SUCCESS(f'成功导入 {len(verb_objs)} 条记录到 {model.__name__} 表中'))

                        # 创建 VerbConjugation 对象
                        conjugation_objs = []
                        for verb in model.objects.all():
                            # 添加 p1, p2 两种分词
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='present',
                                person='p1',
                                conjugated_form=verb.p1
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='present',
                                person='p2',
                                conjugated_form=verb.p2
                            ))

                            # 现在时变位 a1 ~ a6
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='present',
                                person='a1',
                                conjugated_form=verb.a1
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='present',
                                person='a2',
                                conjugated_form=verb.a2
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='present',
                                person='a3',
                                conjugated_form=verb.a3
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='present',
                                person='a4',
                                conjugated_form=verb.a4
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='present',
                                person='a5',
                                conjugated_form=verb.a5
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='present',
                                person='a6',
                                conjugated_form=verb.a6
                            ))

                            # 过去时变位 b1 ~ b6
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='past',
                                person='b1',
                                conjugated_form=verb.b1
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='past',
                                person='b2',
                                conjugated_form=verb.b2
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='past',
                                person='b3',
                                conjugated_form=verb.b3
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='past',
                                person='b4',
                                conjugated_form=verb.b4
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='past',
                                person='b5',
                                conjugated_form=verb.b5
                            ))
                            conjugation_objs.append(VerbConjugation(
                                verb=verb,
                                form_type='past',
                                person='b6',
                                conjugated_form=verb.b6
                            ))

                        # 批量创建 VerbConjugation 对象
                        VerbConjugation.objects.bulk_create(conjugation_objs)
                        self.stdout.write(self.style.SUCCESS(f'成功导入 {len(conjugation_objs)} 条记录到 VerbConjugation 表中'))

                    else:
                        # 清空现有数据
                        model.objects.all().delete()
                        self.stdout.write(f'已清空 {model.__name__} 表中的数据')

                        # 将DataFrame转换为字典列表
                        records = df.to_dict(orient='records')

                        # 处理布尔字段
                        if model == Noun:
                            for record in records:
                                record['m'] = bool(record['m'])
                                # 'gender' 字段已经是字符串，可以直接存储

                        # 批量创建对象
                        objs = [model(**record) for record in records]
                        model.objects.bulk_create(objs)

                        self.stdout.write(self.style.SUCCESS(f'成功导入 {len(objs)} 条记录到 {model.__name__} 表中'))

            except Exception as e:
                self.stdout.write(self.style.ERROR(f'导入 {file_name} 时发生错误: {str(e)}'))
